import os
import shutil
from fastapi import File, UploadFile, Form, Body, Query
from typing import List
from langchain.text_splitter import CharacterTextSplitter
from langchain.vectorstores import FAISS
from langchain.document_loaders import TextLoader,PyPDFLoader,Docx2txtLoader
from langchain.embeddings import HuggingFaceEmbeddings
from langchain_core.documents import Document
import os
import configs
from server.utils import APIResponse





# 创建知识库
def create_kb(
        kb_name: str = Body(...,description="知识库名称"),
        embedding_model: str = Body(...,description="嵌入模型")
):
    return _create_kb(kb_name,embedding_model)
def _create_kb(kb_name,embedding_model):
    model_name = configs.MODEL_PATH+embedding_model
    print(f'模型路径{model_name}')
    print(os.getcwd())
    model_kwargs = {'device': 'cpu'}
    encode_kwargs = {'normalize_embeddings': True}
    embeddings = HuggingFaceEmbeddings(model_name=model_name,
                                       model_kwargs=model_kwargs,
                                       encode_kwargs=encode_kwargs)


    tmp_doc = Document(page_content='',metadata={})
    db = FAISS.from_documents([tmp_doc],embedding=embeddings)
    # db.docstore._dict = {}   该方法删除会引发bug
    db.delete([db.index_to_docstore_id[0]])

    # db.index.ntotal       # 索引中的文档数量

    db.save_local('./KB/' + kb_name)
    print(f"📚知识库： {kb_name} 已创建")
    return APIResponse(code=200, msg=f"📚知识库： {kb_name} 已创建")



# 删除知识库
def delete_kb(kb_name: str = Form(...,description="知识库名称")):
    return _delete_kb(kb_name)
def _delete_kb(kb_name: str = Form(...,description="知识库名称")):
    shutil.rmtree(f"./KB/{kb_name}") # 删除文件夹
    print(f"📚知识库： {kb_name} 已删除")
    return APIResponse(code=200, msg=f"📚知识库： {kb_name} 已删除")



# 更新知识库
def update_kb(
        files: List[UploadFile] = File(..., description="上传文件"),
        kb_name: str = Form(...,description="知识库名称", examples=["examples"]),
        override: bool = Form(...,description="是否覆盖已有知识库"),
        chunk_size: int = Form(...,description="分句大小"),
        chunk_overlap: int = Form(...,description="分句重叠"),
        split_str: str = Form(...,description="分隔符"),
        embedding_model: str = Form(...,description="嵌入模型"),
):
    return _update_kb(files,kb_name,override,chunk_size,chunk_overlap,split_str,embedding_model)
def _update_kb(
        files,kb_name,override,chunk_size,chunk_overlap,split_str,embedding_model):
    model_name = configs.MODEL_PATH+embedding_model
    model_kwargs = {'device': 'cpu'}
    encode_kwargs = {'normalize_embeddings': True}
    embeddings = HuggingFaceEmbeddings(model_name=model_name,
                                       model_kwargs=model_kwargs,
                                       encode_kwargs=encode_kwargs)
    # 加载知识库
    db = FAISS.load_local(f"./KB/{kb_name}",embeddings,allow_dangerous_deserialization = True)
    # 读取文件名类型
    file_type = os.path.splitext(files)[1]
    # 文件读取
    if file_type == '.txt':
        loader = TextLoader(files,encoding='utf8')
    if file_type ==  '.docx':
        loader = Docx2txtLoader(files)
    if file_type == '.pdf':
        loader == PyPDFLoader(files)
    documents = loader.load()
    # 切割
    text_splitter = CharacterTextSplitter(chunk_size=chunk_size,
                                          chunk_overlap=chunk_overlap,
                                          separator=split_str)
    docs = text_splitter.split_documents(documents)
    # 展示存入的数据
    # for index,doc in enumerate(docs):
    #     print(f'切片索引：{index}')
    #     print(doc.page_content)
    #     print('---------------------------------------------------------')
    # 更新知识库
    db.add_documents(docs)
    # 查看当前db知识库的片段

    # 去重(暂不更新)
    print(db.docstore._dict)
    # 知识库保存
    db.save_local('./KB/'+kb_name)
    print(f"📚知识库： {kb_name} 已更新")
    return APIResponse(code=200, msg=f"📚知识库： {kb_name} 已更新")



# 查询现有知识库
def show_kbs():
    return _show_kbs()
def _show_kbs():
    kb_path = './KB'
    entries = os.listdir(kb_path)
    folders = [f for f in entries if os.path.isdir(os.path.join(kb_path, f))]
    if len(folders)== 0:
        print('当前尚未有 📚知识库')
    else:
        print('现有 📚知识库 如下：')
        for kb_name in folders:
            print(kb_name)
    return APIResponse(code=200, msg=f"📚知识库 已全部获取")


def load_kb(kb_name: str = Form(...,description="知识库名称"),
            embedding_model: str = Form(..., description="嵌入模型"),
            ):
    return _load_kb(kb_name,embedding_model)
def _load_kb(kb_name,embedding_model):
    model_name = configs.MODEL_PATH + embedding_model
    model_kwargs = {'device': 'cpu'}
    encode_kwargs = {'normalize_embeddings': True}
    embeddings = HuggingFaceEmbeddings(model_name=model_name,
                                       model_kwargs=model_kwargs,
                                       encode_kwargs=encode_kwargs)
    # 加载知识库
    db = FAISS.load_local(f"./KB/{kb_name}", embeddings, allow_dangerous_deserialization=True)
    configs.CURRENT_KB = db
    return APIResponse(code=200, msg=f"📚知识库 已切换到： {kb_name}")